import sys, os, time, gc
import numpy as np
from coverages.utils import get_layer_outs_new, percent_str, percent
from tensorflow.keras import backend as K
from tensorflow.keras.models import model_from_json
from coverages.coverage import AbstractCoverage


class Mutation_Kill(AbstractCoverage):
    def __init__(self, model, params):
        self.pre_neural_num = int(params.pre_neural_num)
        self.num_class = int(params.num_class)
        mutation_path = "D:\Projects\PycharmProjects\Grad_/output/mutant/" + params.model + "/" + params.dataset
        json_path = "D:\Projects\PycharmProjects\Grad_/data/neural_networks/" + params.dataset + "_" + params.model + ".json"
        self.mutation_models = os.path.join(mutation_path, params.approach)
        self.filelist = []
        for file in sorted(os.listdir(self.mutation_models)):
            self.filelist.append(file)
        self.model = model
        self.temp_result = []
        self.result_label = []
        json_file = open(json_path, 'r')
        loaded_model_json = json_file.read()
        json_file.close()
        self.muta_model = model_from_json(loaded_model_json)
        self.reward_table = [[0 for j in range(self.pre_neural_num * 5)] for i in range(self.num_class)]

    def get_measure_state(self):
        return self.reward_table

    def set_measure_state(self, state):
        self.reward_table = state

    def reset_measure_state(self):
        self.reward_table = [[0 for j in range(self.pre_neural_num * 5)] for i in range(self.num_class)]

    # 杀死率
    def get_current_coverage(self, with_implicit_reward=False):
        return percent(np.array(self.reward_table).sum(), self.pre_neural_num * 5 * self.num_class)

    def test(self, test_inputs, test_labels, with_implicit_reward=False):
        i = 0
        # self.total = len(test_inputs)
        test_labels = self.model.predict(test_inputs)
        test_labels_arg = np.argmax(test_labels, axis=1)
        start_time = time.time()
        for f in self.filelist:
            # print(f)
            # print("mutation_model_" + str(i))
            index1 = int(i / (self.pre_neural_num * 5))
            index2 = int(i % (self.pre_neural_num * 5))
            i = i + 1
            if self.reward_table[index1][index2] == 1:
                continue
            self.muta_model.load_weights(os.path.join(self.mutation_models, f))
            self.temp_result = self.muta_model.predict(test_inputs)
            self.result_label = np.argmax(self.temp_result, axis=1)
            # print(str((abs(self.result_label - test_labels_arg)).sum()))
            for result, label in zip(self.result_label, test_labels_arg):
                if result != label:
                    self.reward_table[index1][index2] = 1
                    break
            K.clear_session()
        print("time: " + str(time.time() - start_time))
        return percent(np.array(self.reward_table).sum(), self.pre_neural_num * 5 * self.num_class)
